/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "gtest/gtest.h"

#include "runtime/rt_model.h"
#include "external/runtime/rt_error_codes.h"
#include "graph/ge_tensor.h"
#include "graph/op_desc.h"
#include "graph/compute_graph.h"
#include "graph/utils/op_desc_utils.h"
#include "graph/utils/attr_utils.h"
#include "graph/utils/tensor_utils.h"
#include "graph/utils/graph_utils.h"
#include "graph/compute_graph.h"
#include "common/util/op_info_util.h"
#include "../fe_test_utils.h"

#define protected public
#define private public
#include "adapter/tbe_adapter/tbe_task_builder_adapter.h"
#include "adapter/common/op_store_adapter_manager.h"
#include "adapter/tbe_adapter/tbe_op_store_adapter.h"
#include "fusion_manager/fusion_manager.h"
#include "task_builder/task_builder.h"

#undef private
#undef protected

using namespace std;
using namespace testing;
using namespace ge;
using namespace fe;

FEOpsStoreInfo taskBuilderUnitTbeOpinfoAdapter  {
        6,
        "tbe-builtin",
        EN_IMPL_HW_TBE,
        "./air/test/engines/nneng/ut/testcase/fusion_engine/ops_kernel_store/fe_config/tbe_opinfo",
        "",
        false,
        false
};

std::vector<FEOpsStoreInfo> taskBuilderUnitAapter{
        taskBuilderUnitTbeOpinfoAdapter ,
};

#define SET_SIZE 128

class UTEST_TaskBuilder : public testing::Test
{
protected:
    static void SetOpDecSize(NodePtr& node){
        OpDesc::Vistor<GeTensorDesc> tensors = node->GetOpDesc()->GetAllInputsDesc();
        for (int i = 0; i < node->GetOpDesc()->GetAllInputsDesc().size(); i++){
            ge::GeTensorDesc tensor = node->GetOpDesc()->GetAllInputsDesc().at(i);
            ge::TensorUtils::SetSize(tensor, SET_SIZE);
            node->GetOpDesc()->UpdateInputDesc(i, tensor);
        }
        OpDesc::Vistor<GeTensorDesc> tensorsOutput = node->GetOpDesc()->GetAllOutputsDesc();
        for (int i = 0; i < tensorsOutput.size(); i++){
            ge::GeTensorDesc tensorOutput = tensorsOutput.at(i);
            ge::TensorUtils::SetSize(tensorOutput, SET_SIZE);
            node->GetOpDesc()->UpdateOutputDesc(i, tensorOutput);
        }
    }
    void SetUp()
    {
        rtContext_t rtContext;
        assert(rtCtxCreate(&rtContext, RT_CTX_GEN_MODE, 0) == ACL_RT_SUCCESS);
        assert(rtCtxSetCurrent(rtContext) == ACL_RT_SUCCESS);


        node_ = CreateNode();
        context_ = CreateContext();

        FusionManager::Instance(AI_CORE_NAME).op_store_adapter_manager_ = make_shared<fe::OpStoreAdapterManager>();
        task_builder_ = shared_ptr<TaskBuilder> (new (nothrow) TaskBuilder());
        task_builder_->task_defs_ = &task_defs_;

        TbeOpStoreAdapter tbeAdapter;
        std:: map<string, string> options;
        FusionManager::Instance(AI_CORE_NAME).op_store_adapter_manager_ = make_shared<fe::OpStoreAdapterManager>();
        Configuration::Instance(fe::AI_CORE_NAME).ops_store_info_vector_ = (taskBuilderUnitAapter);
        FusionManager::Instance(fe::AI_CORE_NAME).op_store_adapter_manager_->Initialize(options, fe::AI_CORE_NAME);
    }

    void TearDown()
    {
        task_builder_.reset();
        DestroyContext(context_);
        node_.reset();
        task_defs_.clear();

        rtContext_t rtContext;
        assert(rtCtxGetCurrent(&rtContext) == ACL_RT_SUCCESS);
        assert(rtCtxDestroy(rtContext) == ACL_RT_SUCCESS);


    }

    static NodePtr CreateNode()
    {
        FeTestOpDescBuilder builder;
        builder.SetName("test_tvm");
        builder.SetType("conv");
        builder.SetInputs({1});
        builder.SetOutputs({1});
        builder.AddInputDesc({1,1,1,1}, ge::FORMAT_NCHW, ge::DT_FLOAT);
        builder.AddOutputDesc({1,1,1,1}, ge::FORMAT_NCHW, ge::DT_FLOAT);
        auto node = builder.Finish();

        const char tbeBin[] = "tbe_bin";
        vector<char> buffer(tbeBin, tbeBin+strlen(tbeBin));
        OpKernelBinPtr tbeKernelPtr = std::make_shared<OpKernelBin>("test_tvm", std::move(buffer));
        node->GetOpDesc()->SetExtAttr(OP_EXTATTR_NAME_TBE_KERNEL, tbeKernelPtr);

        ge::AttrUtils::SetInt(node->GetOpDesc(), "_fe_imply_type", (int64_t)EN_IMPL_CUSTOM_TBE);
        ge::AttrUtils::SetStr(node->GetOpDesc(), "tvm_magic", "RT_DEV_BINARY_MAGIC_ELF");
        ge::AttrUtils::SetBool(node->GetOpDesc(), "is_first_node", true);
        ge::AttrUtils::SetBool(node->GetOpDesc(), "is_last_node", true);
        ge::AttrUtils::SetStr(node->GetOpDesc(), node->GetOpDesc()->GetName()+"_kernelname", "kernelname");

        SetOpDecSize(node);

        return node;
    }

    static RunContext CreateContext()
    {
        rtStream_t stream = nullptr;
        rtModel_t model = nullptr;

        assert(rtStreamCreate(&stream, 0) == ACL_RT_SUCCESS);
        assert(rtModelCreate(&model, 0) == ACL_RT_SUCCESS);
        assert(rtModelBindStream(model, stream, 0) == ACL_RT_SUCCESS);

        RunContext context;
        context.model = model;
        context.stream = stream;
        context.dataMemSize = 100;
        context.dataMemBase = (uint8_t *) (intptr_t) 1000;
        context.weightMemSize = 200;
        context.weightMemBase = (uint8_t *) (intptr_t) 1100;
        context.weightsBuffer = Buffer(20);

        return context;
    }

    static void DestroyContext(RunContext& context)
    {
        assert(rtModelUnbindStream(context.model, context.stream) == ACL_RT_SUCCESS);
        assert(rtModelDestroy (context.model) == ACL_RT_SUCCESS);
        assert(rtStreamDestroy (context.stream) == ACL_RT_SUCCESS);
    }

    static void ReleaseGlobalResouce() {
        for (auto stub_func : kStubFuncs) {
            delete[] stub_func;
        }
        kStubFuncs.clear();
    }

protected:
    NodePtr node_ { nullptr };
    RunContext context_;
    std::vector<domi::TaskDef> task_defs_;
    std::shared_ptr<TaskBuilder> task_builder_;
    std::shared_ptr<fe::OpStoreAdapterManager> op_store_adapter_manager_;
};

TEST_F(UTEST_TaskBuilder, destroyhandle_success)
{
    ccHandle_t handle = nullptr;
    Status ret = CreateHandle(&handle);
    EXPECT_EQ(ret, fe::SUCCESS);
    ret = DestroyHandle(&handle);
    EXPECT_EQ(ret, fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_TaskGenCallback_null_error)
{
    EXPECT_NE(TaskBuilder::TaskGenCallback(nullptr, nullptr), fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_TaskGenCallback_taskBuilderMap_invalid)
{
    rtTaskInfo_t task_info;

    EXPECT_NE(TaskBuilder::TaskGenCallback(context_.model, &task_info), fe::SUCCESS);

    TaskBuilder::task_builder_map_[context_.model] = nullptr;
    EXPECT_NE(TaskBuilder::TaskGenCallback(context_.model, &task_info), fe::SUCCESS);
    TaskBuilder::task_builder_map_.erase(context_.model);
}


TEST_F(UTEST_TaskBuilder, case_TaskGenCallback_task_kernel_stubFunc_error)
{
    rtTaskInfo_t task_info;
    task_info.type = RT_MODEL_TASK_KERNEL;
    rtKernelTaskInfo_t& kernel_task = task_info.u.kernelTask;
    kernel_task.stubFunc = nullptr;

    TaskBuilder::task_builder_map_[context_.model] = task_builder_.get();
    EXPECT_NE(TaskBuilder::TaskGenCallback(context_.model, &task_info), fe::SUCCESS);
    TaskBuilder::task_builder_map_.erase(context_.model);
}

TEST_F(UTEST_TaskBuilder, case_TaskGenCallback_task_event_record_success)
{
    rtTaskInfo_t task_info;
    task_info.type = RT_MODEL_TASK_EVENT_RECORD;

    TaskBuilder::task_builder_map_[context_.model] = task_builder_.get();
    EXPECT_EQ(TaskBuilder::TaskGenCallback(context_.model, &task_info), fe::SUCCESS);
    TaskBuilder::task_builder_map_.erase(context_.model);
}

TEST_F(UTEST_TaskBuilder, case_TaskGenCallback_task_kernel_ex_success)
{
    rtTaskInfo_t task_info;
    task_info.type = RT_MODEL_TASK_KERNEL_EX;

    int args;
    task_info.u.kernelTaskEx.argsSize = sizeof(args);
    task_info.u.kernelTaskEx.args = (uint8_t*)&args;

    task_builder_->orig_op_indexes_.push_back(0);

    TaskBuilder::task_builder_map_[context_.model] = task_builder_.get();
    EXPECT_EQ(TaskBuilder::TaskGenCallback(context_.model, &task_info), fe::SUCCESS);
    TaskBuilder::task_builder_map_.erase(context_.model);
}

TEST_F(UTEST_TaskBuilder, case_TaskGenCallback_task_kernel_ex_origOpIndexes_error)
{
    rtTaskInfo_t task_info;
    task_info.type = RT_MODEL_TASK_KERNEL_EX;

    task_builder_->orig_op_indexes_.clear();

    TaskBuilder::task_builder_map_[context_.model] = task_builder_.get();
    EXPECT_NE(TaskBuilder::TaskGenCallback(context_.model, &task_info), fe::SUCCESS);
    TaskBuilder::task_builder_map_.erase(context_.model);
}

TEST_F(UTEST_TaskBuilder, case_TaskGenCallback_task_fusion_start_success)
{
    rtTaskInfo_t task_info;
    task_info.type = RT_MODEL_TASK_FUSION_START;

    TaskBuilder::task_builder_map_[context_.model] = task_builder_.get();
    EXPECT_EQ(TaskBuilder::TaskGenCallback(context_.model, &task_info), fe::SUCCESS);
    TaskBuilder::task_builder_map_.erase(context_.model);
}

TEST_F(UTEST_TaskBuilder, case_FillKernelContextMultiKernel_error)
{
    vector<uint32_t> orig_op_indexes = {0, 1, 2};
    rtAllKernelTaskInfo_t kernel_task;

    uint16_t args_offset[] = {0, sizeof(uint64_t)};
    kernel_task.argsOffset = args_offset;
    kernel_task.argsCount = 2;

    ccOpContext op_context;
    op_context.kernelType = ccKernelType::CCE_AI_CORE;
    op_context.isFlowtable = true;

    domi::KernelContext kernel_context;
    EXPECT_EQ(TaskBuilder::FillKernelContextMultiKernel(orig_op_indexes, kernel_task, op_context, kernel_context), fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_FillKernelContextMultiKernel_cce_ai_core)
{
    vector<uint32_t> orig_op_indexes = {1};
    rtAllKernelTaskInfo_t kernel_task;

    uint16_t args_offset[] = {0, sizeof(uint64_t)};
    kernel_task.argsOffset = args_offset;
    kernel_task.argsCount = 2;

    ccOpContext op_context;
    op_context.kernelType = ccKernelType::CCE_AI_CORE;
    op_context.isFlowtable = true;

    domi::KernelContext kernel_context;
    EXPECT_EQ(TaskBuilder::FillKernelContextMultiKernel(orig_op_indexes, kernel_task, op_context, kernel_context), fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_FillKernelContextMultiKernel_cce_te)
{
    vector<uint32_t> orig_op_indexes = {1, 2, 3};
    rtAllKernelTaskInfo_t kernel_task;

    uint16_t args_offset[] = {0, sizeof(uint64_t)};
    kernel_task.argsOffset = args_offset;
    kernel_task.argsCount = 2;

    ccOpContext op_context;
    op_context.kernelType = ccKernelType::TE_AI_CORE;
    op_context.isFlowtable = true;

    domi::KernelContext kernel_context;
    EXPECT_EQ(TaskBuilder::FillKernelContextMultiKernel(orig_op_indexes, kernel_task, op_context, kernel_context), fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_FillKernelContextMultiKernel_te_failed)
{
    vector<uint32_t> orig_op_indexes = {};
    rtAllKernelTaskInfo_t kernel_task;

    uint16_t args_offset[] = {1, sizeof(uint64_t)};
    kernel_task.argsOffset = args_offset;
    kernel_task.argsCount = 2;

    ccOpContext op_context;
    op_context.kernelType = ccKernelType::TE_AI_CORE;
    op_context.isFlowtable = true;

    domi::KernelContext kernel_context;
    EXPECT_NE(TaskBuilder::FillKernelContextMultiKernel(orig_op_indexes, kernel_task, op_context, kernel_context), fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_FillKernelContext_kernel_type_te_ai_core)
{
    vector<uint32_t> orig_op_indexes = {1};
    rtKernelTaskInfo_t kernel_task;

    uint16_t args_offset[] = {0, sizeof(uint64_t)};
    kernel_task.argsOffset = args_offset;
    kernel_task.argsCount = 2;

    ccOpContext op_context;
    op_context.kernelType = ccKernelType::TE_AI_CORE;

    domi::KernelContext kernel_context;
    EXPECT_EQ(TaskBuilder::FillKernelContext(orig_op_indexes, kernel_task, op_context, kernel_context), fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_FillFlowtable_success)
{
    domi::KernelContext kernel_context;
    rtKernelTaskInfo_t kernel_task;
    domi::KernelDef kernel_def;

    kernel_context.set_is_flowtable(true);
    kernel_task.argsCount = 2;

    uint64_t value;
    uint8_t args[sizeof(uint64_t) + sizeof(uint32_t)];
    *(uint64_t *)args = (intptr_t)&value;
    *(uint32_t *)(args + sizeof(uint64_t)) = sizeof(value);
    kernel_task.args = (uint8_t*)&args;

    uint16_t args_offset[] = {0, sizeof(uint64_t)};
    kernel_task.argsOffset = args_offset;

    EXPECT_EQ(TaskBuilder::FillFlowtable(kernel_context, kernel_task, kernel_def), fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_FillFlowtable_kernelTask_argsCount_error)
{
    domi::KernelContext kernel_context;
    rtKernelTaskInfo_t kernel_task;
    domi::KernelDef kernel_def;

  kernel_context.set_is_flowtable(true);
    kernel_task.argsCount = 1;

    EXPECT_NE(TaskBuilder::FillFlowtable(kernel_context, kernel_task, kernel_def), fe::SUCCESS);
}

TEST_F(UTEST_TaskBuilder, case_FillFlowtable_flowtableSize_error)
{
    domi::KernelContext kernel_context;
    rtKernelTaskInfo_t kernel_task;
    domi::KernelDef kernel_def;

    kernel_context.set_is_flowtable(true);
    kernel_task.argsCount = 2;

    uint32_t n = 0;
    kernel_task.args = (uint8_t*)&n;

    uint16_t args_offset[] = {0, 0};
    kernel_task.argsOffset = args_offset;

    EXPECT_NE(TaskBuilder::FillFlowtable(kernel_context, kernel_task, kernel_def), fe::SUCCESS);
}
